package symbolTable;

import gen.GoParser;
import gen.GoParserBaseListener;

import org.antlr.v4.runtime.Token;
import org.antlr.v4.runtime.tree.ParseTreeProperty;
import org.antlr.v4.runtime.tree.ParseTreeWalker;
import org.antlr.v4.runtime.tree.TerminalNode;
import typeCheck.TypeChecker;

import java.util.LinkedHashMap;
import java.util.List;

/**
 * The symbol table class consists of a map from the symbol name to the symbol itself,
 * and a structure that holds the mapping relationship between the node and the action of the node.
 * It is mainly responsible for collecting symbols in the program, including variable symbols, function symbols, and type symbols,
 * when traversing the syntax tree for the first time, and saving their scopes for later work.
 */

public class SymbolTable extends GoParserBaseListener {

    private LinkedHashMap<String, Symbol> symbols = new LinkedHashMap<String, Symbol>();

    private GlobalScope globalScope;

    private Scope currentScope;

    private ParseTreeProperty<Scope> scopes = new ParseTreeProperty<Scope>();

    /**
     * get symbols
     * @return symbols
     */
    public LinkedHashMap<String, Symbol> getSymbols() {
        return this.symbols;
    }

    /**
     * get scopes
     * @return scopes
     */
    public ParseTreeProperty<Scope> getScopes() {
        return this.scopes;
    }

    /**
     * get globalScope
     * @return globalScope
     */
    public GlobalScope getGlobalScope() {
        return this.globalScope;
    }

    /**
     * get the type_Context inside ()
     * @param ctx type_Context
     * @return type_Context without ()
     */
    public static GoParser.Type_Context getEmbeddedType_(GoParser.Type_Context ctx) {
        if (ctx.type_() != null)
            return getEmbeddedType_(ctx.type_());
        return ctx;
    }

    /**
     * get the nonNamedTypeContext inside ()
     * @param ctx nonNamedTypeContext
     * @return nonNamedTypeContext without ()
     */
    public static GoParser.NonNamedTypeContext getEmbeddedNonNamedType(GoParser.NonNamedTypeContext ctx) {
        if (ctx.nonNamedType() != null)
            return getEmbeddedNonNamedType(ctx.nonNamedType());
        return ctx;
    }

    /**
     * Define the symbols in the parameters into the function scope
     *
     * @param ctx        ParametersContext
     * @param func    functionSymbol to define
     */
    private void parseParameters(GoParser.ParametersContext ctx, FunctionSymbol func) {
        // get all parameter declaration
        List<GoParser.ParameterDeclContext> paramCtxs = ctx.parameterDecl();

        // Process parameter declaration separately
        for (GoParser.ParameterDeclContext paramCtx : paramCtxs) {
            List<TerminalNode> params = paramCtx.IDENTIFIER();  // get list of name

            // process every parameter in parameter declaration
            for (TerminalNode param : params) {
                String name = param.getText();
                 VariableSymbol var = new VariableSymbol(name);
                 var.scope = func;

                 // if parameters is sub-node of resultContext, just declare result for function or method
                 if (ctx.getParent() instanceof GoParser.ResultContext) {
                     func.results.put(name, var);
                 } else {
                     defineSymbol(var, param.getSymbol());
                 }
            }
        }
    }

    /**
     * Used to generate sub-scope when a function is assigned to a variable
     * @param block  functionBody
     */
    private void parseBlock(GoParser.BlockContext block) {
        // walk subtree manually
        ParseTreeWalker walker = new ParseTreeWalker();
        walker.walk(this, block);
    }

    /**
     * Defines a symbol at the current scope
     *
     * @param symbol symbol
     * @param token token of symbol
     */
    private void defineSymbol(Symbol symbol, Token token) {
        // check whether current scope has same-name symbol
        if (this.currentScope.search(symbol.name) == null) {
            this.currentScope.define(symbol);
            this.symbols.put(symbol.name, symbol);
        } else {
            TypeChecker.error(token, symbol.name + " has been declared");
        }
    }

    /**
     * declare vars with expressions
     * @param exprs expressions
     * @param vars  vars
     */
    private void parseVarsSpecWithExpressions(List<GoParser.ExpressionContext> exprs, List<TerminalNode> vars) {
        int varIndex = 0;
        for (GoParser.ExpressionContext expr : exprs) {
            if (expr.primaryExpr() != null
            && expr.primaryExpr().operand() != null
            && expr.primaryExpr().operand().literal() != null
            && expr.primaryExpr().operand().literal().functionLit() != null) {
                // create a function symbol
                FunctionSymbol func = new FunctionSymbol(vars.get(varIndex).getText(), this.currentScope);

                // define symbol and scope
                defineSymbol(func, vars.get(varIndex++).getSymbol());
                this.currentScope = func;

                // define parameters and result
                parseParameters(expr.primaryExpr().operand().literal().functionLit().signature().parameters(), func);
                if (expr.primaryExpr().operand().literal().functionLit().signature().result() != null) {
                    if (expr.primaryExpr().operand().literal().functionLit().signature().result().parameters() != null) {
                        parseParameters(expr.primaryExpr().operand().literal().functionLit().signature().result().parameters(), func);
                    }
                }

                // create a local scope and define function body
                this.currentScope = new LocalScope(this.currentScope);
                this.scopes.put(expr.primaryExpr().operand().literal().functionLit().block(), this.currentScope);
                parseBlock(expr.primaryExpr().operand().literal().functionLit().block());
                this.currentScope = this.currentScope.getEnclosingScope().getEnclosingScope();
            } else {
                // create a variable symbol
                String name = vars.get(varIndex).getText();
                VariableSymbol variableSymbol = new VariableSymbol(name);
                variableSymbol.scope = this.currentScope;
                defineSymbol(variableSymbol, vars.get(varIndex++).getSymbol());
            }
        }
    }

    @Override
    public void enterSourceFile(GoParser.SourceFileContext ctx) {
        // initiate globalScope and currentScope
        this.globalScope = new GlobalScope(null);
        this.currentScope = this.globalScope;
    }

    @Override
    public void exitSourceFile(GoParser.SourceFileContext ctx) {
        // print the information of symbols in symbol table
        for (String name : this.symbols.keySet()) {
            System.out.println(name + "   " + this.symbols.get(name).type);
        }
    }

    @Override
    public void enterFunctionDecl(GoParser.FunctionDeclContext ctx) {
        // get function information
        String name = ctx.IDENTIFIER().getText();
        FunctionSymbol func = new FunctionSymbol(name, this.currentScope);

        // save symbol and scope
        defineSymbol(func, ctx.IDENTIFIER().getSymbol());
        this.currentScope = func;
        this.scopes.put(ctx, this.currentScope);

        // define parameters and result of function
        parseParameters(ctx.signature().parameters(), func);
        if (ctx.signature().result() != null) {
            if (ctx.signature().result().parameters() != null) {
                parseParameters(ctx.signature().result().parameters(), func);
            } else {
                func.results.put("res", new VariableSymbol("res"));
            }
        }

        // create a local scope and define function body
        this.currentScope = new LocalScope(this.currentScope);
        this.scopes.put(ctx.block(), this.currentScope);
        parseBlock(ctx.block());
        this.currentScope = this.currentScope.getEnclosingScope();
    }

    @Override
    public void exitFunctionDecl(GoParser.FunctionDeclContext ctx) {
        this.currentScope = this.currentScope.getEnclosingScope();
    }

    @Override
    public void enterMethodDecl(GoParser.MethodDeclContext ctx) {
        // get method information
        String name = ctx.IDENTIFIER().getText();
        FunctionSymbol func = new FunctionSymbol(name, this.currentScope);

        // save symbol and scope
        defineSymbol(func, ctx.IDENTIFIER().getSymbol());
        this.currentScope = func;
        this.scopes.put(ctx, this.currentScope);

        // define receiver, parameters and result of method
        parseParameters(ctx.receiver().parameters(), func);
        parseParameters(ctx.signature().parameters(), func);
        if (ctx.signature().result() != null
        && ctx.signature().result().parameters() != null) {
            parseParameters(ctx.signature().result().parameters(), func);
        }

        // create a local scope and define function body
        this.currentScope = new LocalScope(this.currentScope);
        this.scopes.put(ctx.block(), this.currentScope);
        parseBlock(ctx.block());
        this.currentScope = this.currentScope.getEnclosingScope();
    }

    @Override
    public void exitMethodDecl(GoParser.MethodDeclContext ctx) {
        this.currentScope = this.currentScope.getEnclosingScope();
    }

    @Override
    public void enterConstSpec(GoParser.ConstSpecContext ctx) {
        // Save all the leaf nodes representing variables to the list
        List<TerminalNode> vars = ctx.IDENTIFIER();

        // declare constVar
        for (TerminalNode var : vars) {
            String name = var.getText();
            VariableSymbol constVar = new VariableSymbol(name);
            constVar.scope = this.currentScope;
            defineSymbol(constVar, var.getSymbol());
        }
    }

    @Override
    public void enterVarSpec(GoParser.VarSpecContext ctx) {
        // Save all the leaf nodes representing variables to the list
        List<TerminalNode> vars = ctx.IDENTIFIER();

        if (ctx.expressionList() != null) {
            List<GoParser.ExpressionContext> exprs = ctx.expressionList().expression();

            // parse vars with expressions
            parseVarsSpecWithExpressions(exprs, vars);
        } else if(ctx.type_() != null
        && getEmbeddedType_(ctx.type_()).typeLit() != null
        && getEmbeddedType_(ctx.type_()).typeLit().functionType() != null) {

            // all vars declare as function
            for (TerminalNode func : vars) {
                // get information of function symbol
                String name = func.getText();
                FunctionSymbol funcSymbol = new FunctionSymbol(name, this.currentScope);

                // define symbol and scope
                defineSymbol(funcSymbol, func.getSymbol());
                this.currentScope = funcSymbol;
                this.scopes.put(func, this.currentScope);

                // define parameters and result of function symbol
                parseParameters(getEmbeddedType_(ctx.type_()).typeLit().functionType().signature().parameters(), funcSymbol);
                if (getEmbeddedType_(ctx.type_()).typeLit().functionType().signature().result() != null
                && getEmbeddedType_(ctx.type_()).typeLit().functionType().signature().result().parameters() != null) {
                    parseParameters(getEmbeddedType_(ctx.type_()).typeLit().functionType().signature().result().parameters(), funcSymbol);
                }
                this.currentScope = this.currentScope.getEnclosingScope();
            }
        } else {
            // all vars declare as variable
            for (TerminalNode var : vars) {

                // get information of symbol
                String name = var.getText();
                VariableSymbol varSymbol = new VariableSymbol(name);
                varSymbol.scope = this.currentScope;

                // define symbol
                defineSymbol(varSymbol, var.getSymbol());
            }
        }
    }

    @Override
    public void enterShortVarDecl(GoParser.ShortVarDeclContext ctx) {
        // Save all the leaf nodes representing variables and expressions to the list
        List<TerminalNode> vars = ctx.IDENTIFIER();
        List<GoParser.ExpressionContext> exprs = ctx.expressionList().expression();

        // parse vars with expressions
        parseVarsSpecWithExpressions(exprs, vars);
    }
}
